import pandas as pd
import numpy as np
import pydicom
from tqdm import tqdm
import os
from einops import rearrange

import torch
from torch.utils.data import Dataset, DataLoader, Subset

from torchvision.transforms.v2 import Compose,Resize, Normalize, ColorJitter, RandomAffine, RandomErasing, ToTensor,GaussianBlur
from train_224_224 import TransformerMode



test_pids_path = r'train_test_data/datasets/sample_submission.csv'
test_path = r'train_test_data/datasets/test_images'
test_pids = pd.read_csv(test_pids_path)['patient_id'].values

#datasets 
def fetch_img_paths(train_img_path):
    img_paths = []
    print('Scanning directories...')
    for patient in tqdm(os.listdir(train_img_path)):
        for scan in os.listdir(os.path.join(train_img_path, patient)):
            scan_path=os.listdir(os.path.join(train_img_path, patient, scan))
            scans = [
                [filename[:-4], os.path.join(train_img_path, patient, scan, filename)] \
                for filename in scan_path\
                if filename.endswith('.dcm')  and 10<=int(filename[:-4])<=1500
                    ]
            scans=sorted(scans,key=lambda x:int(x[0]))
            img_path=np.array(scans)
            img_paths.append(img_path[:,1].tolist())
    return img_paths


def select_elements_with_spacing(input_list): #读取40张图片
    length = len(input_list)
    if length <= 40:
        return input_list
    else:
        indexs = np.linspace(0, length - 1, 40).astype(int)
    return [input_list[index] for index in indexs]


def preprocess_jpeg(dicm_img_path):
    dicom=clip_rescale_extract_dicom_image(pydicom.read_file(dicm_img_path))
    dicom=np.array(dicom)
    dicoms=np.expand_dims(dicom,axis=2)
    dicoms=np.concatenate((dicoms,dicoms,dicoms),axis=-1)
    return dicoms


def clip_rescale_extract_dicom_image(dicom_ds):
    image = dicom_ds.pixel_array
    
    # find rescale params
    if ("RescaleIntercept" in dicom_ds) and ("RescaleSlope" in dicom_ds):
        intercept = float(dicom_ds.RescaleIntercept)
        slope = float(dicom_ds.RescaleSlope)

    # find clipping params
    center = int(dicom_ds.WindowCenter)
    width = int(dicom_ds.WindowWidth)
    low = center - width / 2  
    high = center + width / 2    

    image = (image * slope) + intercept
    image = np.clip(image, low, high)
    image = (image / np.max(image) * 255).astype(np.int16)
    return image

# dataset
class AbdominalData(Dataset):  # TODO 考虑添加分割的图像
    def __init__(self):
        super().__init__()
        # collect all the image instance paths
        self.img_paths = fetch_img_paths(test_path)
        self.transform = Compose([
            Resize((224, 224), antialias=True),# 进行重新训练  先用小的分辨率训练再使用大的分辨率训练
            ColorJitter(brightness=0.2),  # Randomly adjust brightness
            ColorJitter(contrast=0.2),  # Randomly adjust contrast
            GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)), #添加高斯模糊
            RandomAffine(degrees=0, scale=(0.8, 1.2)),
            RandomErasing(p=0.2, scale=(0.02, 0.2)),  # Coarse dropout
            Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]),
            ToTensor(),
        ])

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        dicom_images = select_elements_with_spacing(self.img_paths[idx])
        patient_id = dicom_images[0].split('/')[-3]
        images = []

        for d in dicom_images:
            image = preprocess_jpeg(d)
            images.append(image)
        images = np.array(images)/1.0
        image = torch.tensor(images, dtype=torch.float)
        image = rearrange(image, 't  h w c ->   t c h w')
        image = self.transform(image)
        image = rearrange(image, ' t c h w ->  c t h w')
        return image,patient_id

#datasets loader
data=DataLoader(AbdominalData(),batch_size=1,shuffle=False,num_workers=0)

# model = torch.load('model.pth')
model=torch.load('checkpoints/0.647_model_180.pth')
model=model.cuda()
model.eval()

test_pids = []
pred_list = []

for dcoms, pid in data: 
    pid=pid[0]
    print(pid)
    if pid not in test_pids:
        test_pids.append(pid)
        with torch.no_grad():
            dcoms = dcoms.cuda()
            output = model(dcoms)[:, :-1]
            pred = torch.sigmoid(output).data.cpu().numpy().round(3)
            pred_list.append(pred)

pred_df = pd.DataFrame(np.vstack(pred_list))
pred_df.columns = ['bowel_healthy', 'bowel_injury', 'extravasation_healthy',
       'extravasation_injury', 'kidney_healthy', 'kidney_low', 'kidney_high',
       'liver_healthy', 'liver_low', 'liver_high', 'spleen_healthy',
       'spleen_low', 'spleen_high']

pred_df['patient_id'] = test_pids[:] # TODO 生成代码

pred_df[['patient_id', 'bowel_healthy', 'bowel_injury', 'extravasation_healthy',
       'extravasation_injury', 'kidney_healthy', 'kidney_low', 'kidney_high',
       'liver_healthy', 'liver_low', 'liver_high', 'spleen_healthy',
       'spleen_low', 'spleen_high']].to_csv('submission.csv', index=None)


# !head /kaggle/working/submission.csv



